import matplotlib.pyplot as plt
import numpy as np

def plot_map(img, map):
    plt.figure(), plt.imshow(img,cmap='gray')
    plt.scatter(map[..., 0], map[..., 1], c='green', marker='*')

def plot_tensor(tensor):
    # 如果 Tensor 的值不在 [0, 1] 范围内，需要进行归一化
    image_tensor = (tensor - tensor.min()) / (tensor.max() - tensor.min())

    # 将通道维度移到最后，因为 matplotlib 要求图像数据的形状为 (H, W, C)
    rows = 2
    cols = 2
    fig, axes = plt.subplots(rows, cols, figsize=(10, 10))
    if image_tensor.ndim > 3:
        for channel in range(image_tensor.size(1)):
            img = image_tensor[:, channel, ...]
            img = img.squeeze()
            img = img.permute(1,0)
            img_np = img.detach().numpy()

            row = channel // cols
            col = channel % cols
            axes[row,col].imshow(img_np)
            axes[row,col].axis('off')

    image_tensor = image_tensor.permute(1, 2, 0)

    image_np = image_tensor.numpy()

    plt.imshow(image_np)
    plt.axis('off')
    plt.show()
